import tushare as ts
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import TimeSeriesSplit
import backtrader as bt
import matplotlib.pyplot as plt

import pandas as pd
import tushare as ts
from sklearn.preprocessing import StandardScaler

from sklearn.svm import SVR, SVC

ts.set_token('b91f8e85973449023ab471c313af892fed91df1fab7fc669abc3de8f')
pro = ts.pro_api()

# 获取日线数据（按日期升序排列）
df = pro.daily(
    ts_code='601398.SH',
    start_date='20100101',
    end_date='20241231',
    fields='trade_date,open,high,low,close,vol'
).sort_values('trade_date')  # 确保升序

# 获取复权因子（按日期升序排列）
adj_df = pro.adj_factor(
    ts_code='601398.SH',
    start_date='20100101',
    end_date='20241231'
).sort_values('trade_date')

# 合并数据
df = df.merge(adj_df[['trade_date', 'adj_factor']], on='trade_date', how='left')

# 计算前复权价格（基准为最新日期）
latest_adj_factor = df['adj_factor'].iloc[-1]  # 正确获取最新因子
for col in ['open', 'high', 'low', 'close']:
    df[col] = df[col] * df['adj_factor'] / latest_adj_factor

# 调整成交量
df['vol'] = df['vol'] / df['adj_factor']

# 格式处理
df['trade_date'] = pd.to_datetime(df['trade_date'])
df.set_index('trade_date', inplace=True)
df.sort_index(ascending=True, inplace=True)
df.rename(columns={'vol': 'volume'}, inplace=True)

# 删除冗余列
df.drop(columns=['adj_factor'], inplace=True)

# 打印前复权数据示例
print(df.head())


# ================== 特征工程 ==================
def calculate_features(data):
    """计算技术指标特征（优化版）"""
    # 价格指标
    data['MA5'] = data['close'].rolling(5).mean()
    data['MA20'] = data['close'].rolling(20).mean()
    data['MA_diff'] = data['MA5'] - data['MA20']

    # 波动率指标
    data['ATR'] = data['high'].rolling(14).max() - data['low'].rolling(14).min()

    # 成交量指标
    data['VOL_MA10'] = data['volume'].rolling(10).mean()
    data['VOL_MA20'] = data['volume'].rolling(20).mean()

    # 动量指标（优化RSI计算）
    delta = data['close'].diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)
    avg_gain = gain.rolling(14).mean()
    avg_loss = loss.rolling(14).mean()
    data['RSI'] = 100 - (100 / (1 + (avg_gain / avg_loss.replace(0, np.nan))))

    # 滞后特征（过去3天收盘价）
    for lag in [1, 2, 3]:
        data[f'close_lag{lag}'] = data['close'].shift(lag)

    return data.dropna()


feature_df = calculate_features(df.copy())

# ================== 标签生成 ==================
# 定义标签：次日上涨为1，下跌为0
feature_df['target'] = (feature_df['close'].shift(-1) > feature_df['close']).astype(int)
feature_df.dropna(inplace=True)  # 删除最后一行无标签数据

features = ['MA5', 'MA20', 'MA_diff', 'ATR', 'VOL_MA10', 'VOL_MA20', 'RSI',
            'close_lag1', 'close_lag2', 'close_lag3']

tss = TimeSeriesSplit(n_splits=50)
predictions = []
scaler = StandardScaler()  # 新增标准化器

for train_idx, test_idx in tss.split(feature_df):
    # 时间序列分割
    X_train = feature_df.iloc[train_idx][features]
    y_train = feature_df.iloc[train_idx]['target']
    X_test = feature_df.iloc[test_idx][features]

    # 修改点2：添加数据标准化
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 修改点3：替换为SVM模型
    model = SVC(
        kernel='rbf',        # 高斯核
        C=1.0,              # 正则化参数
        gamma='scale',      # 自动计算gamma
        probability=True,   # 启用概率预测
        random_state=42
    )
    model.fit(X_train_scaled, y_train)

    # 修改点4：使用标准化后的数据预测
    pred = model.predict(X_test_scaled)
    pred_proba = model.predict_proba(X_test_scaled)[:, 1]  # 获取正类概率

    test_dates = feature_df.iloc[test_idx].index
    for date, (pred_val, proba) in zip(test_dates, zip(pred, pred_proba)):
        print(f"日期: {date.date()}, 信号: {pred_val}, 置信度: {proba:.2%}")

    predictions.append(pd.Series(pred, index=test_dates))

feature_df['signal'] = pd.concat(predictions)


# ================== Backtrader回测 ==================
# 自定义数据类（解决AttributeError关键）
class SignalData(bt.feeds.PandasData):
    lines = ('signal',)  # 新增信号线
    params = (
        ('datetime', None),
        ('open', 'open'),
        ('high', 'high'),
        ('low', 'low'),
        ('close', 'close'),
        ('volume', 'volume'),
        ('signal', 'signal'),
    )


class MLStrategy(bt.Strategy):
    params = (
        ('printlog', True),
        ('stake', 1000),  # 默认固定数量
        ('risk_pct', 0),  # 0表示不使用百分比模式
    )

    def __init__(self):
        self.signal = self.data.lines.signal
        self.order = None

        # 动态选择仓位控制器
        if self.params.risk_pct > 0:
            self.sizer = bt.sizers.PercentSizer()
            self.sizer.params.percents = self.params.risk_pct * 100

    def next(self):
        if self.order:
            return

        # 多头开仓
        if self.signal[0] == 1 and not self.position:
            # 根据模式选择仓位大小
            if self.params.risk_pct > 0:
                self.order = self.buy()  # 由sizer自动计算数量
            else:
                self.order = self.buy(size=self.params.stake)

        # 平仓
        elif self.signal[0] == 0 and self.position:
            self.order = self.close()

    # 订单通知和日志方法保持不变

    def notify_order(self, order):
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'BUY EXECUTED, Price: {order.executed.price:.2f}')
            elif order.issell():
                self.log(f'SELL EXECUTED, Price: {order.executed.price:.2f}')
            self.order = None

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()}, {txt}')


# ================== 回测引擎配置 ==================
cerebro = bt.Cerebro()

# 数据预处理
numeric_cols = ['open', 'high', 'low', 'close', 'volume', 'signal']
feature_df[numeric_cols] = feature_df[numeric_cols].apply(pd.to_numeric, errors='coerce')
feature_df = feature_df.dropna(subset=numeric_cols)
feature_df['signal'] = feature_df['signal'].astype(int)

# 加载数据
data = SignalData(dataname=feature_df)
cerebro.adddata(data)

# 策略参数
cerebro.addstrategy(MLStrategy)
cerebro.broker.setcash(1000000)
cerebro.broker.setcommission(commission=0.001)

# 添加分析器
cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe')
cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')

# ================== 运行回测 ==================
print('初始资金: %.2f' % cerebro.broker.getvalue())
results = cerebro.run()
print('最终资金: %.2f' % cerebro.broker.getvalue())

# 输出分析结果
strat = results[0]
print('夏普比率:', strat.analyzers.sharpe.get_analysis())
print('最大回撤:', strat.analyzers.drawdown.get_analysis()['max']['drawdown'])

# 可视化
cerebro.plot(style='candlestick')